/*
 * Description: 任务栏卡死监听
 */
#include "taskbar_monitor.h"
#include <sstream>
#include <tlhelp32.h>
#include <psapi.h>
#include "named_pipe_util.h"
#include "proc_util.h"
#include "string_util.h"

#pragma comment(lib, "Psapi.lib")
#pragma comment(lib, "Advapi32.lib")

namespace {
    constexpr int MONITOR_INTERVAL_S = 300;
    const std::string SHELL_TRAYWND = "Shell_TrayWnd";
    const std::string TASKBAR_WND = "TaskbarWnd";
    const std::string DESK_TOP_WND = "DeskTopWnd";
    const std::string STARTS_WND_OF_TASKBAR = "StartsWndOfTaskbar";
    const std::string NOTIFICATION_AERA_WND_OF_TASKBAR = "NotificationAreaWndOfTaskbar";
    const std::string PROCESS_WND_OF_TASKBAR = "ProcessWndOfTaskbar";
    constexpr UINT32 TASKBAR_WINDOW = 1; // 任务栏窗口
    constexpr UINT32 DESKTOP_WINDOW = 2; // 桌面窗口

    std::map<std::string const, UINT32> g_stuckWindow = {
    {TASKBAR_WND, TASKBAR_WINDOW},
    {STARTS_WND_OF_TASKBAR, TASKBAR_WINDOW},
    {NOTIFICATION_AERA_WND_OF_TASKBAR, TASKBAR_WINDOW},
    {PROCESS_WND_OF_TASKBAR, TASKBAR_WINDOW},
    {DESK_TOP_WND, DESKTOP_WINDOW},
    };

    constexpr UINT32 MONITOR_PERIOD_SECOND = 5;
    constexpr int PROCESS_NOEXIST = 4;
    constexpr int APP_HANG = 2;
    constexpr int WND_HANG = 1;
    constexpr int INIT_VALUE = 0;
    constexpr DWORD WND_HANG_MSG_TIMEOUT = 5000; // 检测窗口是否卡死超时时间5s
    constexpr UINT32 MONITOR_CNT_COMM_THRESHOLD = 3;

    const std::string RE_BAR_WINDOWS = "ReBarWindow32";
    const std::string MS_TASK_WINDOWS = "MSTaskSwWClass";
    const std::string TRAY_NOTIFY_WND = "TrayNotifyWnd";
    const std::string SYS_PAGER = "SysPager";
    const std::string ROOL_BAR_WINDOWS = "ToolbarWindow32";
    const std::string MSTASK_LIST_WCLASS = "MSTaskListWClass";
    const std::string SHELLDELL_DEF_VIEW = "SHELLDLL_DefView";
    const std::string START = "Start";
    const std::string SYS_LIST_VIEW = "SysListView32";

    std::map<WCT_OBJECT_STATUS, std::string> g_objectStatus = {
        // 线程状态 拒绝访问
        {WctStatusNoAccess, "WctStatusNoAccess"},
        // 线程状态 运行中
        {WctStatusRunning, "WctStatusRunning"},
        // 线程状态 被阻塞
        {WctStatusBlocked, "WctStatusBlocked"},
        // 线程状态 PidOnly
        {WctStatusPidOnly, "WctStatusPidOnly"},
        // 线程状态 是线程独有的
        {WctStatusPidOnlyRpcss, "WctStatusPidOnlyRpcss"},
        // 分发对象状态 被拥有
        {WctStatusOwned, "WctStatusOwned"},
        // 分发对象状态 未被拥有
        {WctStatusNotOwned, "WctStatusNotOwned"},
        // 分发对象状态 被遗弃
        {WctStatusAbandoned, "WctStatusAbandoned"},
        // 所有对象 未知状态
        {WctStatusUnknown, "WctStatusUnknown"},
        // 所有对象 错误
        {WctStatusError, "WctStatusError"},
        {WctStatusMax, "WctStatusMax"}
    };

    constexpr UINT32 CLASS_NAME_LEN = 256;
}

TaskBarMonitor& TaskBarMonitor::GetInstance()
{
    static TaskBarMonitor instance;
    return instance;
}

void TaskBarMonitor::Init()
{
    if (taskbarStuckMonitorThread.joinable()) {
        taskbarStuckMonitorThread.join();
    }
    taskbarStuckMonitorThread = std::thread([this]() {
        std::unique_lock<std::mutex> lck(mtx);
        while (!stopMonitorFlag) {
            while (!cv.wait_for(lck, std::chrono::seconds(MONITOR_INTERVAL_S), [this]() {
                return this == nullptr || stopMonitorFlag;
                })) {
                MonitorTaskBarStuck();
            }
        }
        });
}

void TaskBarMonitor::DeInit()
{
    stopMonitorFlag = true;
    cv.notify_all();
    if (taskbarStuckMonitorThread.joinable()) {
        taskbarStuckMonitorThread.join();
    }
}

void TaskBarMonitor::MonitorTaskBarStuck()
{
    InitTaskBarStuckEventParams();
    if (!IsTaskBarStuck()) {
        return;
    }
    if ((windowsClassOfStuck != "") && (g_stuckWindow.count(windowsClassOfStuck) != 0)) {
        // 任务栏卡死，获取卡死弹窗类型
        windowsTypeOfStuck = g_stuckWindow[windowsClassOfStuck];
    }
    GetExplorerWaitChain();
    LOG_INFO("TaskBar Stuck");
    NotifyTaskBarStuck2Agent();
}

void TaskBarMonitor::NotifyTaskBarStuck2Agent()
{
    // 通知agent_service
    std::ostringstream dataStream;
    dataStream << "{\"StuckTime\":\"";
    dataStream << StringUtil::FormatDateTime("%Y-%m-%d %H:%M:%S", stuckTime);
    dataStream << "\",";
    dataStream << "\"StuckType\":";
    dataStream << stuckType;
    dataStream << ",";
    dataStream << "\"WaitChain\":\"";
    dataStream << waitChain;
    dataStream << "\",";
    dataStream << "\"WindowsTypeOfStuck\":";
    dataStream << windowsTypeOfStuck;
    dataStream << ",";
    dataStream << "\"WindowsClassOfStuck\":\"";
    dataStream << windowsClassOfStuck;
    dataStream << "\"}";
    std::string taskBarStuckMsg = "TaskBarStuck:" + dataStream.str();
    LOG_INFO("TaskBarStuck msg: {}", taskBarStuckMsg);
    NamedPipeUtil::GetInstance().Send(taskBarStuckMsg);
}


void TaskBarMonitor::InitTaskBarStuckEventParams()
{
    stuckTime = time(nullptr);
    stuckType = INIT_VALUE;
    waitChain = "";
    windowsClassOfStuck = "";
    windowsTypeOfStuck = INIT_VALUE;
}

bool TaskBarMonitor::ProcessHwndNotRespond(const DWORD time, const HWND hwnd)
{
    LRESULT result = SendMessageTimeout(hwnd, WM_NULL, 0, 0, SMTO_BLOCK, time, 0);
    // SendMessageTimeout: 如果函数调用成功，返回非零值,
    // 如果函数调用失败，或超时，返回值是零
    if (result == 0 && GetLastError() == ERROR_TIMEOUT) {
        LOG_INFO("ProcessHwndNotRespond is hang");
        return true;
    }
    return false;
}

bool TaskBarMonitor::ExplorerWndNotRespond(const DWORD time, std::string& hwndName) // 任务栏及其子窗口的检测
{
    std::map<std::string, HWND> hwndMap;
    HWND hwnd = FindWindow(SHELL_TRAYWND.c_str(), NULL); // 获取任务栏窗口句柄
    if (hwnd != nullptr) {
        hwndMap.insert({ TASKBAR_WND, hwnd });
        HWND tmpHwnd = GetStartsWndOfTaskbar(hwnd);
        if (tmpHwnd) {
            hwndMap.insert({ STARTS_WND_OF_TASKBAR, tmpHwnd }); // 获取任务栏子窗口“开始”按钮的窗口句柄
        }
        tmpHwnd = GetNotificationAreaWndOfTaskbar(hwnd);
        if (tmpHwnd) {
            hwndMap.insert({ NOTIFICATION_AERA_WND_OF_TASKBAR, tmpHwnd }); // 获取任务栏子窗口“用户提示通知区域”的窗口句柄
        }
        tmpHwnd = GetProcessWndOfTaskbar(hwnd);
        if (tmpHwnd) {
            hwndMap.insert({ PROCESS_WND_OF_TASKBAR, tmpHwnd }); // 获取任务栏子窗口“应用图标”的窗口句柄
        }
    }
    HWND deskHwnd = nullptr;
    GetDeskTopWindows(deskHwnd);
    if (deskHwnd) {
        hwndMap.insert({ DESK_TOP_WND, deskHwnd }); // 获取桌面窗口句柄
    }

    hwndName = ""; // 初始化窗口
    for (auto it : hwndMap) {
        /* 检测窗口是否有响应 */
        if (ProcessHwndNotRespond(time, it.second)) {
            hwndName = it.first;
            LOG_DEBUG("ProcessHwndNotRespond hwndName: {}", hwndName.c_str());
            return true;
        }
    }
    return false;
}

HWND TaskBarMonitor::GetprocessWndHandle(const std::string& wndName)
{
    HWND hWnd = FindWindow(wndName.c_str(), nullptr);
    if (!hWnd) {
        LOG_ERR("FindWindow {} Fail, err: {}", wndName.c_str(), GetLastError());
        return nullptr;
    }
    return hWnd;
}

bool TaskBarMonitor::IsTaskbarProcessExist()
{
    if (!ProcUtil::IsProcessExists("explorer.exe")) { // 进程不存在了
        LOG_INFO(" explorer.exe is no exist");
        stuckType = PROCESS_NOEXIST; // 进程不存在
        return false;
    } else { // 进程存在，判断是不是仅文件夹打开了，而任务栏仍然不在
        std::vector<DWORD> tokenVec = {};
        std::vector<std::string> explorerCommandLine = {};
        ProcUtil::GetProcessIdByProcessName("explorer.exe", tokenVec);
        int num = 0;
        for (auto it : tokenVec) {
            HANDLE processHandle = OpenProcess(PROCESS_ALL_ACCESS, 0, it);
            if (processHandle == nullptr) {
                LOG_ERR("OpenProcess failed, err:{}", GetLastError());
                continue;
            }
            std::string commandLine = ProcUtil::GetProcessCommandLine(processHandle);
            CloseHandle(processHandle);
            explorerCommandLine.push_back(commandLine);
            LOG_DEBUG("explorer.exe:{}", commandLine.c_str());
            if (commandLine.find("factory") != std::string::npos) {
                num++;
            }
        }
        if (num == tokenVec.size()) {
            stuckType = PROCESS_NOEXIST; // 进程不存在
            return false;
        }
    }
    return true;
}

bool TaskBarMonitor::IsTaskBarStuck()
{
    int count = 0;
    while ((!stopMonitorFlag) && (count < MONITOR_CNT_COMM_THRESHOLD)) { // 连续3次检测到5s无响应，则认为任务栏卡死了
        std::string hwndName = "";
        if (ExplorerWndNotRespond(WND_HANG_MSG_TIMEOUT, hwndName)) { // 窗口无响应
            windowsClassOfStuck = hwndName; // 储存卡死的窗口名
            time_t winHangTime = time(nullptr);
            winHangVec.push_back(winHangTime);
            count++;
            continue;
        } else if (GetprocessWndHandle(SHELL_TRAYWND.c_str()) == NULL) { // 获取不到任务栏窗口句柄
            LOG_INFO("Get exploer hwnd failed.");
            if (!IsTaskbarProcessExist()) { // 检查任务栏进程是否存在
                windowsClassOfStuck = TASKBAR_WND;
                return true;
            }
            winHangVec.clear();
            return false;
        }
        winHangVec.clear(); // 窗口有响应
        return false;
    }
    if (winHangVec.size() < MONITOR_CNT_COMM_THRESHOLD) {
        LOG_DEBUG("TaskBar no hang");
        return false;
    }
    stuckType = WND_HANG;
    LOG_INFO("TaskBar hang");
    return true;
}

HWND TaskBarMonitor::GetProcessWndOfTaskbar(HWND& handle) // 获取“应用图标”的窗口句柄
{
    LOG_INFO("GetProcessWndOfTaskbar enter");
    HWND processWnd = FindWindowEx(handle, 0, RE_BAR_WINDOWS.c_str(), nullptr);
    if (!processWnd) {
        return processWnd;
    }
    processWnd = FindWindowEx(processWnd, 0, MS_TASK_WINDOWS.c_str(), nullptr);
    if (!processWnd) {
        return processWnd;
    }
    processWnd = FindWindowEx(processWnd, 0, MSTASK_LIST_WCLASS.c_str(), nullptr);
    if (!processWnd) {
        return processWnd;
    }
    LOG_INFO("Got MSTaskListWClass hwnd");
    return processWnd;
}

HWND TaskBarMonitor::GetNotificationAreaWndOfTaskbar(HWND& handle) // 获取“用户提示通知区域”窗口句柄
{
    LOG_INFO("GetNotificationAreaWndOfTaskbar enter");
    HWND areaWnd = FindWindowEx(handle, 0, TRAY_NOTIFY_WND.c_str(), nullptr);
    if (!areaWnd) {
        return areaWnd;
    }
    LOG_INFO("Got TrayNotifyWnd hwnd");
    areaWnd = FindWindowEx(areaWnd, 0, SYS_PAGER.c_str(), nullptr);
    if (!areaWnd) {
        return areaWnd;
    }
    LOG_INFO("Got SysPager hwnd");
    areaWnd = FindWindowEx(areaWnd, 0, ROOL_BAR_WINDOWS.c_str(), nullptr);
    if (!areaWnd) {
        return areaWnd;
    }
    LOG_INFO("Got ToolbarWindow32 hwnd");
    return areaWnd;
}

HWND TaskBarMonitor::GetStartsWndOfTaskbar(HWND& handle) // 获取“开始”窗口句柄
{
    LOG_INFO("GetStartsWndOfTaskbar enter");
    HWND startWnd = FindWindowEx(handle, 0, START.c_str(), nullptr);
    if (handle) {
        LOG_INFO("Got Start hwnd");
    }
    return startWnd;
}

static BOOL CALLBACK EnumDeskTopWindowsProc(HWND hwnd, LPARAM lParam)
{
    HWND* pHwnd = reinterpret_cast<HWND*>(lParam);
    TCHAR className[CLASS_NAME_LEN] = { 0 };
    GetClassName(hwnd, className, sizeof(className) / sizeof(TCHAR)); // 枚举所有窗口类
    if (strcmp(className, "Progman") == 0) { // 找到窗口类“Progman”
        HWND childWnd = FindWindowEx(hwnd, 0, SHELLDELL_DEF_VIEW.c_str(), nullptr);
        if (!childWnd) {
            return true;
        }
        childWnd = FindWindowEx(childWnd, 0, SYS_LIST_VIEW.c_str(), nullptr);
        if (childWnd) {
            (*pHwnd) = childWnd;
            return true;
        }
    }
    return true; // 记录遍历其他窗口
}

void TaskBarMonitor::GetDeskTopWindows(HWND& handle)
{
    EnumWindows(EnumDeskTopWindowsProc, reinterpret_cast<LPARAM>(&handle));
}

void TaskBarMonitor::GetExplorerWaitChain()
{
    std::string classNameTaskMan = "Shell_TrayWnd";
    HWND hWnd = FindWindow(classNameTaskMan.c_str(), nullptr);
    if (hWnd) {
        DWORD processId = 0;
        GetWindowThreadProcessId(hWnd, static_cast<LPDWORD>(&processId));
        std::map<std::string, std::string> getWaitChain = {};
        GetWaitChain(processId, getWaitChain);  // 获取等待链
        for (auto it : getWaitChain) {
            waitChain = it.first;
            waitChain += ",";
            waitChain += it.second;
        }
    }
}

void TaskBarMonitor::GetWaitChain(DWORD procesId, std::map<std::string, std::string>& getWaitChain)
{
    LOG_INFO(" GetWaitChain enter");
    HWCT wctHandle = OpenThreadWaitChainSession(0, NULL);
    if (NULL == wctHandle) {
        LOG_ERR("OpenThreadWaitChainSession failed");
        return;
    }
    if (!GrantDebugPrivilege()) {
        LOG_ERR("Could not enable the debug privilege");
    }

    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, procesId);
    if (!hSnapshot) {
        CloseThreadWaitChainSession(wctHandle);
        return;
    }
    THREADENTRY32 thread;
    thread.dwSize = sizeof(thread);
    if (!Thread32First(hSnapshot, &thread)) {
        CloseHandle(hSnapshot);
        CloseThreadWaitChainSession(wctHandle);
        return;
    }
    do {
        if (thread.th32OwnerProcessID != procesId) {
            continue;
        }
        // Open a handle to this specific thread
        HANDLE hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, thread.th32ThreadID);
        if (hThread) {
            // Check whether the thread is still running
            DWORD exitCode;
            GetExitCodeThread(hThread, &exitCode);
            if ((exitCode == STILL_ACTIVE) && PrintWaitChain(thread.th32ThreadID, wctHandle, getWaitChain)) {
                break;
            }
            CloseHandle(hThread);
        }
    } while (Thread32Next(hSnapshot, &thread) && (!stopMonitorFlag));
    CloseThreadWaitChainSession(wctHandle);
    if (hSnapshot != nullptr) {
        CloseHandle(hSnapshot);
        hSnapshot = nullptr;
    }
}

bool TaskBarMonitor::GrantDebugPrivilege()
{
    LOG_DEBUG("GrantDebugPrivilege GrantDebugPrivilege enter");
    bool fSuccess = false;
    HANDLE hToken = nullptr;
    TOKEN_PRIVILEGES tokenPrivileges;
    do {
        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
            LOG_ERR("GrantDebugPrivilege Could not get the process token");
            break;
        }
        tokenPrivileges.PrivilegeCount = 1;
        if (!LookupPrivilegeValue(nullptr, SE_DEBUG_NAME, &tokenPrivileges.Privileges[0].Luid)) {
            LOG_ERR("GrantDebugPrivilege Couldn't lookup SeDebugPrivilege name");
            break;
        }
        tokenPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        if (!AdjustTokenPrivileges(hToken, FALSE, &tokenPrivileges, sizeof(tokenPrivileges), nullptr, nullptr)) {
            LOG_ERR("GrantDebugPrivilege Could not revoke the debug privilege");
            break;
        }
        fSuccess = true;
    } while (false);
    if (hToken) {
        CloseHandle(hToken);
    }
    return fSuccess;
}

bool TaskBarMonitor::PrintWaitChain(DWORD threadId, HWCT wctHandle, std::map<std::string, std::string>& getWaitChain)
{
    LOG_DEBUG("PrintWaitChain enter");
    WAITCHAIN_NODE_INFO NodeInfoArray[WCT_MAX_NODE_COUNT];
    LOG_DEBUG("threadId: {}", threadId);
    BOOL isCycle = FALSE;
    DWORD count = WCT_MAX_NODE_COUNT;
    // Make a synchronous WCT call to retrieve the wait chain.
    if (!GetThreadWaitChain(wctHandle, 0, WCTP_GETINFO_ALL_FLAGS, threadId, &count, NodeInfoArray, &isCycle)) {
        LOG_ERR("GetThreadWaitChain failed with error: {}", GetLastError());
        return false;
    }
    LOG_ERR("additional nodes: {}", count);
    // Check if the wait chain is too big for the array we passed in.
    for (DWORD i = 0; i < min(count, WCT_MAX_NODE_COUNT); i++) {
        // 如果该节点表示被阻塞的线程。（其下个元素就是等待的原因）
        if (NodeInfoArray[i].ObjectType == WctThreadType) {
            // name:processId:Status
            // 进程名：进程ID：线程名->原因
            DWORD processId = NodeInfoArray[i].ThreadObject.ProcessId;
            std::string objectStatus = "";
            if (g_objectStatus.count(NodeInfoArray[i].ObjectStatus) > 0) {
                objectStatus = g_objectStatus[NodeInfoArray[i].ObjectStatus];
                LOG_DEBUG("PrintWaitChain objectStatus: {}", objectStatus.c_str());
            }
            HANDLE proHandle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processId);
            if (proHandle != NULL) {
                WCHAR szName[MAX_PATH] = { 0 };
                GetModuleBaseNameW(proHandle, NULL, szName, sizeof(szName));
                CloseHandle(proHandle);
                std::wstring processNamew = szName;
                std::string processName = StringUtil::WStringToString(processNamew);
                LOG_DEBUG("PrintWaitChain szName: {}", processName.c_str()); // 进程名
                getWaitChain[processName] = objectStatus;
                return true;
            }
        }
    }
    if (isCycle) {
        LOG_ERR("PrintWaitChain !!!Deadlock!!!");
    }
    return false;
}
